"""
CTPRPCGlobal

definitions of pseudo-constants, exceptions and functions used by
CTPRPCServer.py and CTPRPCClient.py
"""

import md5
import sha
import hmac
import cPickle
import threading
import select
import os

__revision__ = "1.00"


# global pseudo-constants

# Encryption modes
CTP_ENCRYPT_NONE	= 0
CTP_ENCRYPT_DES3	= 1
CTP_ENCRYPT_AES		= 2

# Compression modes
CTP_COMPRESS_NONE	= 0
CTP_COMPRESS_ZLIB	= 1
CTP_COMPRESS_BZ2	= 2

# CTP-RPC commands
CMD_HEARTBEAT		= 0
CMD_INVOKE		= 1
CMD_HANDSHAKE_1		= 2
CMD_HANDSHAKE_2		= 3
CMD_LOGIN		= 4

# Responses
RESP_OK			= 0
RESP_EXCEPTION		= 1
RESP_BADMETHOD		= 2
RESP_BADCOMMAND		= 3
RESP_HANDSHAKE_1	= 4
RESP_HANDSHAKE_2	= 5
RESP_BADLOGIN		= 6

# CTP Packet Flags
CTP_FLAG_COMPRESSED	= 0x01
CTP_FLAG_ENCRYPTED	= 0x02
CTP_FLAG_CONTINUED	= 0x04
CTP_FLAG_COMPRESS_DATA	= 0x08

# Other global settings
RSA_KEY_LENGTH		= 512
DEFAULT_KEY_SIZE	= 16
MAX_PACKET_SIZE		= 2	# MBytes
MAX_FILE_TRANSFER_SIZE	= 1	# MBytes


# *******************************************************
# Exception classes
class CTPRPCError ( Exception ): pass
class BadMethodError ( CTPRPCError ): pass
class InvalidResponseError ( CTPRPCError ): pass
class BadLoginOrPassword ( CTPRPCError ): pass


# MAIN

MAX_PACKET_SIZE *= 1048576
MAX_FILE_TRANSFER_SIZE *= 1048576

# build list of supported encryption/compression types
supported_encryption_types = {
				CTP_ENCRYPT_NONE:	[ '', None ],
				CTP_ENCRYPT_AES:	[ 'AES', None ],
				CTP_ENCRYPT_DES3:	[ 'DES3', None ],
			}

supported_compression_types = {
				CTP_COMPRESS_NONE:	[ '', None ],
				CTP_COMPRESS_BZ2:	[ 'bz2', None ],
				CTP_COMPRESS_ZLIB:	[ 'zlib', None ]
			}

#test if we can support encryption
try:
	import Crypto.PublicKey.RSA
	import Crypto.Util.randpool
	
	RSA = Crypto.PublicKey.RSA
	RANDPOOL = Crypto.Util.randpool.RandomPool()

except ImportError:
	supported_encryption_types = { CTP_COMPRESS_NONE: [ '', None ] }

#test supported encryption algorithms
to_del = []
for type, mod in supported_encryption_types.iteritems():
	if not mod[0]: continue
	try:
		mod[1] = getattr ( __import__ ( 'Crypto.Cipher', globals(), locals(), [ mod[0] ] ), mod[0] )
		if not mod[1].key_size: mod[1].key_size = DEFAULT_KEY_SIZE
	except ImportError:
		to_del.append ( type )
for d in to_del:
	del supported_encryption_types[d]

#test supported compression types
to_del = []
for type, mod in supported_compression_types.iteritems():
	if not mod[0]: continue
	try:
		mod[1] = __import__ ( mod[0] )
	except ImportError:
		to_del.append ( type )
for d in to_del:
	del supported_compression_types[d]

is_encryption_supported = len ( supported_encryption_types ) > 1
is_compression_supported = len ( supported_compression_types ) > 1

# FILE PICKLING/UNPICKLING
file_unpicklers = {}
file_pickle_lock = threading.Lock()


# *******************************************************
# CTP_PRF() - pseudo-random function
def CTP_PRF ( key, data, reqlen ):
	"""
	CTP_PRF ( key, data, reqlen ) -> str

	Pseudo-Random Function used internally by CTP-RPC server and client.
	"""
	
	md5_mac = hmac.new ( key, data, md5 ).digest()
	sha_mac = hmac.new ( key, data, sha ).digest()

	retval = ''
	while reqlen > 0:
		md5_new = hmac.new ( md5_mac + key, sha_mac + data, md5 ).digest()
		sha_new = hmac.new ( sha_mac + key, md5_mac + data, sha ).digest()
		retval += md5_new + sha_new
		md5_mac = md5_new
		sha_mac = sha_new
		reqlen -= 36
		
	return retval[:reqlen]
	

# *******************************************************
# read_bytes() - legge n bytes dal socket
def read_bytes ( sock, n ):
	"""
	read_bytes ( sock, n )

	Funzione di basso livello. Legge esattamente ``n`` byte da un socket.
	"""
	
	buf = ''
        timeout = sock.gettimeout()
        while n > 0:
                count = 0
                while count < timeout:
                        try: rl, wl, xl = select.select ( [ sock ], [], [ sock ], 1 )
                        except: return ''
                        if xl: return ''
                        elif rl: break
                        count += 1
                else:
                        return ''
                r_buf = sock.recv ( n )
                if not r_buf: return ''
                buf += r_buf
                n -= len ( r_buf )
        return buf


# *******************************************************
# read_object() - legge un oggetto da un socket
def read_object ( sock, session_id, hmac_key = '', compress_mod = None, encrypt_key = None ):
	"""
	read_object ( sock, session_id, hmac_key, compress_mod, encrypt_key ) -> obj or None

	Funzione CTP di basso livello per leggere un oggetto da un socket.

	Se la fase di handshake e' terminata effettua anche la decompressione/decifratura/verifica del MAC
	se sono state concordate.
	"""
	
	try:
		ret = ctp_sock_read ( sock, hmac_key, compress_mod, encrypt_key )
		if not ret: return None

		blob, flags = ret

		encrypted = flags & CTP_FLAG_ENCRYPTED
		compressed = flags & CTP_FLAG_COMPRESSED
		compress_data = flags & CTP_FLAG_COMPRESS_DATA

		pickler = CTPPickler ( sock, session_id )

		file_pickle_lock.acquire()
		try: file_unpicklers[session_id] = pickler
		except: pass
		file_pickle_lock.release()

		obj = pickler.unpickle ( blob )
		
		file_pickle_lock.acquire()
		try: del file_unpicklers[session_id]
		except: pass
		file_pickle_lock.release()

		files = pickler.unpickled_files
		cur_file = None

		while len(files):
			
			ret = ctp_sock_read ( sock, hmac_key, compress_mod, encrypt_key )
			if not ret: return None

			buf = ret[0]
			
			fobj = cPickle.loads ( buf )
			
			if fobj[0] == 'FILE':
				if cur_file:
					cur_file.seek ( 0, 0 )
					del files[0]
				if not len(files): return None
				recv_fname = fobj[1]
				fname = files[0][0]
				if recv_fname != fname: return None
				cur_file = files[0][1]
			elif fobj[0] == 'DATA':
				if not cur_file: return None
				cur_file.write ( fobj[1] )
			elif fobj[0] == 'END':
				if cur_file:
					cur_file.seek ( 0, 0 )
					del files[0]
				if len(files): return None
			else:
				return None
		
		return ( obj, encrypted, compressed, compress_data )
	except:
		raise
		return None


# *******************************************************
# ctp_sock_read() - legge dati da un socket con il protocollo CTP
def ctp_sock_read ( sock, hmac_key, compress_mod, encrypt_key ):

	blob = ''
	continued = 1
	while continued:
	
		flags = read_bytes ( sock, 2 )
		if not flags: return None
		flags = int ( flags, 16 )
		size = read_bytes ( sock, 8 )
		if not size: return None
		size = int ( size, 16 )
		buffer = read_bytes ( sock, size )
		if not buffer: return None

		encrypted	= flags & CTP_FLAG_ENCRYPTED
		compressed	= flags & CTP_FLAG_COMPRESSED
		continued	= flags & CTP_FLAG_CONTINUED

		#DECRYPTION
		if encrypted and encrypt_key:
			buffer = encrypt_key.decrypt ( buffer )
			size = int ( buffer[:8], 16 )
			buffer = buffer[8:8 + size]

		#DECOMPRESSION
		if compressed and compress_mod: buffer = compress_mod.decompress ( buffer )

		#MAC VERIFICATION
		buffer, mac = cPickle.loads ( buffer )
		calc_mac = hmac.new ( hmac_key, buffer, md5 ).digest()
		calc_mac += hmac.new ( hmac_key, buffer, sha ).digest()
		if calc_mac != mac: return None

		blob += buffer

	return ( blob, flags )


# *******************************************************
# send_object() - invia un oggetto su un socket
def send_object ( sock, obj, session_id, hmac_key = '', compress_mod = None, encrypt_key = None, block_size = 1, compress_data = 1 ):
	"""
	send_object ( sock, obj, session_id, hmac_key, compress_mod, encrypt_key, block_size )

	Funzione CTP di basso livello per inviare un oggetto su un socket.

	Se la fase di handshake e' terminata effettua anche la compressione/cifratura/apposizione MAC
	se sono state concordate.
	"""

	pickler = CTPPickler ( sock, session_id )
	blob = pickler.pickle ( obj )

	ctp_sock_write ( sock, blob, hmac_key, compress_mod, encrypt_key, block_size, compress_data )

	if not compress_data: compress_mod = None

	for f in pickler.pickled_files:

		#get file size
		f.seek ( 0, 2 )
		size = f.tell()
		f.seek ( 0, 0 )

		buf = cPickle.dumps ( ( 'FILE', f.name ) )
		ctp_sock_write ( sock, buf, hmac_key, compress_mod, encrypt_key, block_size, compress_data )

		while 1:
			buf = f.read ( MAX_FILE_TRANSFER_SIZE )
			if not len(buf): break

			buf = cPickle.dumps ( ( 'DATA', buf, ) )
			ctp_sock_write ( sock, buf, hmac_key, compress_mod, encrypt_key, block_size, compress_data )
			
			if len(buf) < MAX_FILE_TRANSFER_SIZE: break
	
	if len(pickler.pickled_files):
		buf = cPickle.dumps ( ( 'END', ) )
		ctp_sock_write ( sock, buf, hmac_key, compress_mod, encrypt_key, block_size, compress_data )


# *******************************************************
# ctp_sock_write() - scrive dati su un socket con il protocollo CTP
def ctp_sock_write ( sock, blob, hmac_key, compress_mod, encrypt_key, block_size, compress_data = 1 ):

	size = len(blob)
	
	cur_ptr = 0
	while size > 0:

		flags = 0
		if compress_data: flags |= CTP_FLAG_COMPRESS_DATA
		
		if size > MAX_PACKET_SIZE:
			size -= MAX_PACKET_SIZE
			to_send = MAX_PACKET_SIZE
			flags |= CTP_FLAG_CONTINUED
		else:
			to_send = size
			size = 0

		obj_str = blob [ cur_ptr : cur_ptr + to_send ]
		cur_ptr += to_send

		#MAC
		if hmac_key:
			mac = hmac.new ( hmac_key, obj_str, md5 ).digest()
			mac += hmac.new ( hmac_key, obj_str, sha ).digest()
			obj = ( obj_str, mac )
			obj_str = cPickle.dumps ( obj )

		#COMPRESSION
		if compress_mod:
			obj_str = compress_mod.compress ( obj_str )
			flags |= CTP_FLAG_COMPRESSED

		#ENCRYPTION
		if encrypt_key:
			obj_str = "%08X" % len(obj_str) + obj_str
			padlen = block_size - ( len(obj_str) % block_size )
			obj_str += "J" * padlen
			#print "dec: ", obj_str[:8].__repr__(), len(obj_str)
			obj_str = encrypt_key.encrypt ( obj_str )
			#print "enc: ", obj_str[:8].__repr__(), len(obj_str)
			flags |= CTP_FLAG_ENCRYPTED

		sock.sendall ( "%02X%08X" % ( flags, len(obj_str) ) )
		sock.sendall ( obj_str )


# *******************************************************
# CTPTempFile class
class CTPTempFile ( object ):
	def __init__ ( self, name ):
		super ( CTPTempFile, self ).__init__()
		self._file = os.tmpfile()
		self.name = name


	# *******************************************************
	# __del__() - distruttore, chiude il file
	def __del__ ( self ):
		self.close()

	
	# *******************************************************
	# __getattr__() - restituisce gli attributi di file
	def __getattr__ ( self, name ):
		f = self.__dict__.get ( "_file" )
		if not f: raise AttributeError
		return getattr ( f, name )


# *******************************************************
# CTPPickler class
class CTPPickler ( object ):
	def __init__ ( self, sock, key ):
		super ( CTPPickler , self ).__init__()
		self._copy_reg = __import__ ( 'copy_reg' )
		self._copy_reg.pickle ( file, self._pickle_file )
		self._copy_reg.pickle ( CTPTempFile, self._pickle_file )

		self.pickled_files = []
		self.unpickled_files = []

		self._key = key


	# *******************************************************
	# pickle() - effettua un pickle di un oggetto
	def pickle ( self, obj ):
		return cPickle.dumps ( obj )
	
	
	# *******************************************************
	# _pickle_file() - effettua un pickle di un oggetto file
	def _pickle_file ( self, f ):
		self.pickled_files.append ( f )
		return ( unpickle_file, ( self._key, f.name ) )


	# *******************************************************
	# unpickle() - effettua un unpickle di un oggetto
	def unpickle ( self, str ):
		return cPickle.loads ( str )
	
	
# *******************************************************
# unpickle_file() - effettua un unpickle di un oggetto file
def unpickle_file ( key, name ):

	tfile = CTPTempFile ( name )

	file_pickle_lock.acquire()
	try: file_unpicklers[key].unpickled_files.append ( ( name, tfile ) )
	except: pass
	file_pickle_lock.release()

	return tfile

